[feat] Enable NVFP4 for model weight loading daemon - #32398
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
I will review this PR as soon as possible |
| self.module_attrs: Dict[str, Dict[str, Any]] = {} | ||
| # MoE runner backend as resolved during this daemon's model load; the | ||
| # client compares it once its own model is built. | ||
| self.moe_runner_backend: str = "" |
There was a problem hiding this comment.
It was sent to the client side to ensure both sides load the weights using the same backend:
As different MoE backends post-process the weight tensor (e.g., permute, reorder) into different physical layout, mixing backends will lead to shape mismatch or numeric error.
|
Thanks for the work here — the round-trip design and the backend-fingerprint guards look solid. My main concern is architectural: most of the new "generic" machinery is really NVFP4-specific knowledge that has leaked into the transport layer — |
Fully agree with this, the quant method should be responsible for maintaining the states transferring between daemon and client. Will revise it soon. |
|
Updated and refactored. |
Motivation
#27139 enables the daemon mode for model weight loading from CUDA IPC. This PR extends its supporting range further by including NVFP4 models (e.g.,
nvidia/Qwen3-30B-A3B-FP4).Modifications
The design of this PR delegates the post-process procedure of NVFP4 weight loading to daemon. The daemon process captures some of the important attributes like
weights_padding_colsandintermediate_size_per_partitionafter loading the raw weights, and handle over them to the client then. At the client side, it just rebinds them into the related quantized parameters after receiving the CUDA IPC handle.Accuracy Tests
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #30752098039
Latest PR Test (Extra): ❌ Run #30752097972